import os.path

from langchain.chat_models import init_chat_model
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
import bs4
from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import START, StateGraph
from typing_extensions import List, TypedDict
from dotenv import load_dotenv
from langchain_ollama import ChatOllama

load_dotenv(".venv/.env")

# llm = init_chat_model("gpt-4o-mini", model_provider="openai")
llm = ChatOllama(
    model="llama3.2:3b",
    temperature=0
)

def init_vector_store():
    db_path = "./chroma_langchain_db"
    # 建议使用ollama3.2，因为它更小一点
    embeddings = OllamaEmbeddings(model="llama3.2:3b")
    store = Chroma(
        collection_name="example_collection",
        embedding_function=embeddings,
        persist_directory="./chroma_langchain_db",  # Where to save data locally, remove if not necessary
    )

    if not os.path.exists(db_path):
        loader = WebBaseLoader(
            web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
            bs_kwargs=dict(
                parse_only=bs4.SoupStrainer(
                    class_=("post-content", "post-title", "post-header")
                )
            ),
        )
        docs = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        all_splits = text_splitter.split_documents(docs)
        # 索引文本块
        store.add_documents(documents=all_splits)

    return store


vector_store = init_vector_store()

prompt = hub.pull("rlm/rag-prompt")

# define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

# define application steps
def retrieve(state: State):
    retrieved_docs = vector_store.similarity_search(state["question"])
    return {"context": retrieved_docs}

def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt.invoke({"question": state["question"], "context": docs_content})
    response = llm.invoke(messages)

    return {"answer": response.content}

# compile application
graph_builder = StateGraph(State).add_sequence(([retrieve, generate]))
graph_builder.add_edge(START, "retrieve")
graph = graph_builder.compile()

# test
response = graph.invoke({"question": "What is Task Decomposition"})
print(response["answer"])

from PIL import Image

png_image = graph.get_graph().draw_mermaid_png()
with open("mermaid.png", "wb") as f:
    f.write(png_image)
image = Image.open("mermaid.png")
image.show()

